from groundingdino.util.inference import load_model, predict
import cv2
from typing import Tuple
import numpy as np
import torch
import groundingdino.datasets.transforms as T
from PIL import Image
import io
import base64
import json

def load_image(image: Image.Image) -> Tuple[np.array, torch.Tensor]:
    transform = T.Compose(
        [
            T.RandomResize([800], max_size=1333),
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )
    image_np = np.asarray(image)  # Convert PIL image to numpy array
    image_transformed,_ = transform(image,None)  # Apply transformations
    return image_np, image_transformed

def init_context(context):
    context.logger.info("Init context... 0%")
    # model = load_model("groundingdino/config/GroundingDINO_SwinT_OGC.py", "weights/groundingdino_swint_ogc.pth",device="gpu")
    model = load_model("groundingdino/config/GroundingDINO_SwinT_OGC.py", "weights/groundingdino_swint_ogc.pth")
    context.user_data.model = model

    print(torch.cuda.is_available())
    print('DONE!')

    context.logger.info("Init context... 100%")

def handler(context, event):
    context.logger.info("Run YOLOv8 model")
    TEXT_PROMPT = "chair . cat . dog ."
    BOX_THRESHOLD = 0.35
    TEXT_THRESHOLD = 0.25

    # 从事件中解码图片
    data = event.body
    buf = io.BytesIO(base64.b64decode(data["image"]))
    image = Image.open(buf).convert("RGB")

    # 加载和转换图片
    image_source, image = load_image(image)

    # 预测结果
    boxes, logits, phrases = predict(
        model=context.user_data.model,
        image=image,
        caption=TEXT_PROMPT,
        box_threshold=BOX_THRESHOLD,
        text_threshold=TEXT_THRESHOLD
        # device="cpu"
    )

    # 将预测结果转换为可序列化的 JSON 格式
    result = {
        "boxes": boxes.tolist(),  # 将 boxes 转换为列表
        "logits": logits.tolist(),  # 将 logits 转换为列表
        "phrases": phrases  # phrases 通常是字符串列表，已经可序列化
    }

    # 将结果打包为 JSON 字符串
    return context.Response(body=json.dumps(result), headers={},
                            content_type="application/json", status_code=200)